import numpy as np


class RFFTransformer:

    def __init__(self, gamma: float = 1., D: int = 1000) -> None:
        super().__init__()
        self.gamma = gamma
        self.D = D
        self.rng = np.random.default_rng(0)

    def transform(self, x):
        if not hasattr(self, 'weight'):
            self.weight = np.sqrt(2 * self.gamma) * self.rng.standard_normal(
                (x.shape[-1], self.D))
            self.bias = self.rng.random(size=(self.D, )) * 2 * np.pi

        lin = x @ self.weight + self.bias
        return np.sqrt(2 / self.D) * np.cos(lin)
